/*************************************************************************
    > File Name: RSA.cpp
    > Author: ZhaoLulu@BJTU
    > E-mail: 19281121@bjtu.edu.cn
    > Created Time: 2021年11月17日 星期三
 ************************************************************************/
#include "RSA.h"

using namespace std;

void mod_exp(mpz_t result, const mpz_t base, const mpz_t exponent, const mpz_t n) {
    mpz_t f, power;
    mpz_init_set_ui(f, 1);
    mpz_init(power);
    mpz_mod(power, base, n);

    char exp[KEY_LENGTH + 10];
    mpz_get_str(exp, 2, exponent);

    for (int i = (int) strlen(exp) - 1; i >= 0; --i) {
        mpz_mul(f, f, f);
        mpz_mod(f, f, n);
        if (exp[i] == '1') {
            mpz_mul(f, f, power);
            mpz_mod(f, f, n);
        }
    }

    mpz_set(result, f);
}

int miller_rabin(mpz_srcptr n, mpz_srcptr nm1, mpz_ptr x, mpz_ptr y, mpz_srcptr q, unsigned long int k) {
    unsigned long int i;
    mpz_powm(y, x, q, n);
    if (mpz_cmp_ui (y, 1L) == 0 || mpz_cmp(y, nm1) == 0) { return 1; }
    for (i = 1; i < k; i++) {
        mpz_powm_ui(y, y, 2L, n);
        if (mpz_cmp(y, nm1) == 0) { return 1; }
        if (mpz_cmp_ui (y, 1L) == 0) { return 0; }
    }
    return 0;
}

int Miller_Rabin(mpz_srcptr n, int reps) {
    int r;
    mpz_t nm1, nm3, x, y, q;
    unsigned long int k;
    gmp_randstate_t grt;
    int is_prime = 1;

    // Fermat test
    mpz_init(nm1);
    mpz_sub_ui(nm1, n, 1L);
    mpz_init(x);
    mpz_init(y);
    mpz_set_ui(x, 210L);
    mpz_powm(y, x, nm1, n);
    if (mpz_cmp_ui (y, 1L) != 0) {
        mpz_clear(nm1);
        mpz_clear(x);
        mpz_clear(y);
        return 0;
    }

    mpz_init(q);
    k = mpz_scan1(nm1, 0L);
    mpz_tdiv_q_2exp(q, nm1, k);
    mpz_init(nm3);
    mpz_sub_ui(nm3, n, 3L);
    gmp_randinit_default(grt);
    for (r = 0; r < reps && is_prime; r++) {
        mpz_urandomm(x, grt, nm3);
        mpz_add_ui(x, x, 2L);
        is_prime = miller_rabin(n, nm1, x, y, q, k);
    }
    gmp_randclear(grt);
    mpz_clear(nm1);
    mpz_clear(nm3);
    mpz_clear(x);
    mpz_clear(y);
    mpz_clear(q);
    return is_prime;
}

mpz_t *generate_primes() {
    gmp_randstate_t grt;
    gmp_randinit_default(grt);
    gmp_randseed_ui(grt, time(nullptr));

    mpz_t key_p, key_q;
    mpz_init(key_p);
    mpz_init(key_q);

    mpz_urandomb(key_p, grt, KEY_LENGTH / 2);
    mpz_urandomb(key_q, grt, KEY_LENGTH / 2);

    auto *result = new mpz_t[2];
    mpz_init(result[0]);
    mpz_init(result[1]);

    mpz_nextprime(result[0], key_p);
    if (Miller_Rabin(result[0], 25)) {
        auto buf_p = new char[KEY_LENGTH / 2 + 10];
        mpz_get_str(buf_p, BASE, result[0]);
        cout << "p = " << buf_p << endl;
    }
    mpz_nextprime(result[1], key_q);
    if (Miller_Rabin(result[1], 25)) {
        auto buf_q = new char[KEY_LENGTH / 2 + 10];
        mpz_get_str(buf_q, BASE, result[1]);
        cout << "q = " << buf_q << endl;
    }

    mpz_clear(key_p);
    mpz_clear(key_q);

    return result;
}

void generate_d(mpz_t e, mpz_t phi, mpz_t *temp, mpz_t *a, mpz_t *b) {
    mpz_t zero, one, i;
    mpz_init_set_str(zero, "0", 10);
    mpz_init_set_str(one, "1", 10);
    mpz_init(i);
    if (mpz_cmp(zero, phi) == 0) {
        mpz_set(*temp, e);
        mpz_set(*a, one);
        mpz_set(*b, zero);
        return;
    }
    mpz_mod(i, e, phi);
    generate_d(phi, i, temp, b, a);
    mpz_div(i, e, phi);
    mpz_mul(i, i, *a);
    mpz_sub(*b, *b, i);
    mpz_clear(zero);
    mpz_clear(one);
    mpz_clear(i);
}

key_pair *generate_key_pair() {
    mpz_t *primes = generate_primes();

    mpz_t key_n, key_e, key_phi;
    mpz_init(key_n);
    mpz_init_set_ui(key_e, E);
    mpz_init(key_phi);

    mpz_mul(key_n, primes[0], primes[1]);
    mpz_sub_ui(primes[0], primes[0], 1);
    mpz_sub_ui(primes[1], primes[1], 1);
    mpz_mul(key_phi, primes[0], primes[1]);

    mpz_t key_d, temp, x, y;
    mpz_init(key_d);
    mpz_init(temp);
    mpz_init(x);
    mpz_init(y);
    generate_d(key_e, key_phi, &temp, &x, &y);
    mpz_add(key_d, key_phi, x);
    mpz_mod(key_d, key_d, key_phi);

    auto buf_n = new char[KEY_LENGTH + 10];
    auto buf_d = new char[KEY_LENGTH + 10];
    mpz_get_str(buf_n, BASE, key_n);
    mpz_get_str(buf_d, BASE, key_d);
    auto *result = new key_pair;
    result->n = buf_n;
    result->d = buf_d;
    result->e = E;

    mpz_clear(primes[0]);
    mpz_clear(primes[1]);
    mpz_clear(key_n);
    mpz_clear(key_e);
    mpz_clear(key_phi);
    mpz_clear(key_d);
    mpz_clear(temp);
    mpz_clear(x);
    mpz_clear(y);
    delete[]primes;

    return result;
}

char *encrypt(const char *plaintext, const char *key_n, int key_e) {
    mpz_t M, C, n, e;
    mpz_init_set_str(M, plaintext, BASE);
    mpz_init_set_ui(C, 0);
    mpz_init_set_str(n, key_n, BASE);
    mpz_init_set_ui(e, key_e);

    mod_exp(C, M, e, n);

    char *result = new char[KEY_LENGTH + 10];
    mpz_get_str(result, BASE, C);

    return result;
}

char *decrypt(const char *ciphertext, const char *key_n, const char *key_d) {
    mpz_t M, C, n, d;
    mpz_init(M);
    mpz_init_set_str(C, ciphertext, BASE);
    mpz_init_set_str(n, key_n, BASE);
    mpz_init_set_str(d, key_d, BASE);

    mpz_powm(M, C, d, n);

    char *result = new char[KEY_LENGTH + 10];
    mpz_get_str(result, BASE, M);

    return result;
}

int main() {
    key_pair *p = generate_key_pair();

    cout << "n = " << p->n << endl;
    cout << "d = " << p->d << endl;
    cout << "e = " << p->e << endl;

    char buf[KEY_LENGTH + 10];
    cout << "请输入待加密明文数字（二进制长度不超过" << KEY_LENGTH << "）：" << endl;
    cin >> buf;

    cout << "相同密钥条件下加密第1次：" << endl;
    char *ciphertext = encrypt(buf, p->n, p->e);
    cout << "密文为：" << ciphertext << endl;
    cout << "相同密钥条件下加密第2次：" << endl;
    ciphertext = encrypt(buf, p->n, p->e);
    cout << "密文为：" << ciphertext << endl;
    cout << "相同密钥条件下加密第3次：" << endl;
    ciphertext = encrypt(buf, p->n, p->e);
    cout << "密文为：" << ciphertext << endl;
    char *plaintext = decrypt(ciphertext, p->n, p->d);
    cout << "明文为：" << plaintext << endl;

    if (strcmp(buf, plaintext) != 0) {
        cout << "解密失败！" << endl;
    } else {
        cout << "解密成功！" << endl;
    }

    return 0;
}
